import os
import logging
import pickle
import numpy as np
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
import pinecone
from tqdm import tqdm
from pinecone import Pinecone, ServerlessSpec
os.environ["LOKY_MAX_CPU_COUNT"] = "12"
pinecone = Pinecone(api_key="e97e2b32-22d1-4eb0-bbd3-02d95c2484ce")
# 设置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 索引名称
index_name = "mnist-index"

# 获取现有索引列表
existing_indexes = pinecone.list_indexes()

# 检查索引是否存在，如果存在就删除
# 这个if是否需要，要看情况而定
# 比如有的时候，如果不想要重复删除在创建，这个if就可以不要
if any(index['name'] == index_name for index in existing_indexes):
    logging.info(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    logging.info(f"索引 '{index_name}' 已成功删除。")
else:
    logging.info(f"索引 '{index_name}' 不存在，将创建新索引。")

# 创建新索引
logging.info(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
logging.info(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
logging.info(f"已成功连接到索引 '{index_name}'。")
digits = datasets.load_digits()
X = digits.data
y = digits.target

# 将数据集划分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 使用80%的数据创建Pinecone索引
train_size = int(0.8 * len(X_train))
X_train_pinecone = X_train[:train_size]
y_train_pinecone = y_train[:train_size]

# 上传数据到Pinecone
batch_size = 100
for i in tqdm(range(0, len(X_train_pinecone), batch_size), desc="Uploading data to Pinecone"):
    batch = [(str(j), X_train_pinecone[j].tolist()) for j in range(i, min(i + batch_size, len(X_train_pinecone)))]
    index.upsert(batch)

logging.info(f"成功创建索引，并上传了{len(X_train_pinecone)}条数据")

# 训练并测试k=11的KNN模型
k = 11
knn = KNeighborsClassifier(n_neighbors=k)
knn.fit(X_train, y_train)
y_pred = knn.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)

# 打印准确率
logging.info(f"当k={k}时，使用Pinecone的准确率: {accuracy:.4f}")

# 保存最佳KNN模型
with open('best_knn_model.pkl', 'wb') as f:
    pickle.dump(knn, f)